package com.hk.core.data.r2dbc.repository.support;

import com.hk.commons.query.page.QueryPage;
import com.hk.commons.query.page.SimpleQueryPage;
import com.hk.commons.util.CollectionUtils;
import com.hk.commons.util.Lazy;
import com.hk.commons.util.SpringContextHolder;
import com.hk.core.jdbc.SelectArguments;
import com.hk.core.jdbc.SelectStatement;
import io.r2dbc.spi.Parameters;
import io.r2dbc.spi.Row;
import io.r2dbc.spi.RowMetadata;
import lombok.AccessLevel;
import lombok.NoArgsConstructor;
import org.springframework.data.r2dbc.convert.EntityRowMapper;
import org.springframework.data.r2dbc.core.R2dbcEntityOperations;
import org.springframework.data.r2dbc.dialect.DialectResolver;
import org.springframework.data.r2dbc.dialect.R2dbcDialect;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.ArrayList;
import java.util.function.BiFunction;

/**
 * @author Kevin
 */
@NoArgsConstructor(access = AccessLevel.PROTECTED)
public abstract class R2dbcRepositorySupport {

    private static final Lazy<R2dbcEntityOperations> R2DBC_ENTITY_OPERATIONS = Lazy.of(() -> SpringContextHolder.getBean(R2dbcEntityOperations.class));

    private static final Lazy<R2dbcDialect> R2DBC_DIALECT = Lazy.of(() -> DialectResolver.getDialect(R2DBC_ENTITY_OPERATIONS.get()
            .getDatabaseClient().getConnectionFactory()));

    public static <T> Flux<T> queryForList(
            SelectArguments arguments, Class<T> returnType) {
        var selectStatement = SelectStatement.buildSelect(arguments);
        var databaseClient = R2DBC_ENTITY_OPERATIONS.get().getDatabaseClient();
        var selectSpec = databaseClient
                .sql(() -> String.format("%s %s", selectStatement.selectSql().toString(),
                        R2DBC_DIALECT.get().limit().getLimitOffset(arguments.getPageSize(),
                                (long) arguments.getStartRowIndex() * arguments.getPageSize())));
        var parameters = selectStatement.parameters();
        if (CollectionUtils.isNotEmpty(parameters)) {
            for (int i = 0, size = parameters.size(); i < size; i++) {
                selectSpec = selectSpec.bind(i, parameters.get(i));
            }
        }
        return selectSpec.map(mappingFunction(returnType)).all();
    }

    private static <T> BiFunction<Row, RowMetadata, T> mappingFunction(Class<T> returnType) {
        return new EntityRowMapper<>(returnType, R2DBC_ENTITY_OPERATIONS.get().getConverter());
    }

    public static <T> Mono<QueryPage<T>> queryForPage(SelectArguments arguments,
                                                      Class<T> returnType) {
        var selectStatement = SelectStatement.buildSelect(arguments);
        var databaseClient = R2DBC_ENTITY_OPERATIONS.get().getDatabaseClient();
        var countSpec = databaseClient.sql(() -> selectStatement.countSql().toString());
        var parameters = selectStatement.parameters();
        if (CollectionUtils.isNotEmpty(parameters)) {
            for (int i = 0, size = parameters.size(); i < size; i++) {
                countSpec = countSpec.bind(i, Parameters.in(parameters.get(i)));
            }
        }
        return countSpec.map((r, md) -> r.get(0, Long.class))
                .first()
                .defaultIfEmpty(0L)
                .flatMap(count -> {
                    if (count == 0L) {
                        return Mono.just(new SimpleQueryPage<>(new ArrayList<>(), 0L,
                                arguments.getStartRowIndex() + 1, arguments.getPageSize()));
                    }
                    var selectSpec =
                            databaseClient.sql(() -> String.format("%s %s",
                                    selectStatement.selectSql().toString(),
                                    R2DBC_DIALECT.get().limit().getLimitOffset(arguments.getPageSize(),
                                            (long) arguments.getStartRowIndex() * arguments.getPageSize())));
                    if (CollectionUtils.isNotEmpty(parameters)) {
                        for (int i = 0, size = parameters.size(); i < size; i++) {
                            selectSpec = selectSpec.bind(i, Parameters.in(parameters.get(i)));
                        }
                    }
                    return selectSpec.map(mappingFunction(returnType)).all()
                            .collectList()
                            .map(list -> new SimpleQueryPage<>(list, count, arguments.getStartRowIndex() + 1,
                                    arguments.getPageSize()));
                });
    }
}
